/*
 * Copyright 2013-2016 ksyun.com, Inc. or its affiliates. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License").
 * You may not use this file except in compliance with the License.
 * A copy of the License is located at
 *
 *  http://ksyun.com/apache2.0
 *
 * or in the "license" file accompanying this file. This file is distributed
 * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
 * express or implied. See the License for the specific language governing
 * permissions and limitations under the License.
 */
package com.kingsoft.miot.bridge.iam.common;

import org.apache.commons.codec.binary.Hex;

import com.kingsoft.miot.bridge.iam.common.util.SdkHttpUtils;
import com.kingsoft.miot.bridge.iam.common.util.StringUtils;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.UnsupportedEncodingException;
import java.net.URI;
import java.security.DigestInputStream;
import java.security.MessageDigest;
import java.util.*;

import static com.kingsoft.miot.bridge.iam.common.util.StringUtils.UTF8;

/**
 * Signer implementation that signs requests with the AWS4 signing protocol.
 */
public class AWS4Signer {

	public static final String HMAC_SHA256_ALGORITHM = "HmacSHA256";

	public static final String AWS4_TERMINATOR = "aws4_request";

	public static final String AWS4_SIGNING_ALGORITHM = "AWS4-HMAC-SHA256";

	public static final String X_AMZ_DATE = "X-Amz-Date";

	public static final String AUTHORIZATION = "Authorization";

	public static final String HOST = "Host";

	/**
	 * Service name override for use when the endpoint can't be used to determine
	 * the service name.
	 */
	private String serviceName;

	/**
	 * Region name override for use when the endpoint can't be used to determine the
	 * region name.
	 */
	private String regionName;

	/**
	 * The scope of the signature.
	 */
	private final String scope;

	/**
	 * UTC formatted version of the signing time stamp.
	 */
	private final String formattedSigningDateTime;

	/**
	 * UTC Formatted Signing date with time stamp stripped
	 */
	private final String formattedSigningDate;

	/**
	 * Construct a new AWS4 signer instance. By default, enable double url-encoding.
	 */
	public AWS4Signer(String regionName, String serviceName) {
		long signingDateTimeMilli = System.currentTimeMillis();
		this.serviceName = serviceName;
		this.regionName = regionName;
		this.formattedSigningDate = StringUtils.formatDateStamp(signingDateTimeMilli);
		this.formattedSigningDateTime = StringUtils.formatTimestamp(signingDateTimeMilli);
		this.scope = String.format("%s/%s/%s/%s", this.formattedSigningDate, this.regionName, this.serviceName,
				AWS4_TERMINATOR);
	}

	public void sign(SignableRequest<?> request, String ak, String sk) {

		addHostHeader(request);
		request.addHeader(X_AMZ_DATE, this.formattedSigningDateTime);

		String contentSha256 = calculateContentHash(request);
		final String canonicalRequest = createCanonicalRequest(request, contentSha256);
		final String stringToSign = createStringToSign(canonicalRequest);
		final byte[] signingKey = deriveSigningKey(sk);
		final byte[] signature = computeSignature(stringToSign, signingKey);

		request.addHeader(AUTHORIZATION, buildAuthorizationHeader(request, signature, ak));
	}

	private void addHostHeader(SignableRequest<?> request) {
		// AWS4 requires that we sign the Host header so we
		// have to have it in the request by the time we sign.

		final URI endpoint = request.getUri();
		String hostHeader = endpoint.getHost();
		if (SdkHttpUtils.isUsingNonDefaultPort(endpoint)) {
			hostHeader = hostHeader + ":" + endpoint.getPort();
		}
		request.addHeader(HOST, hostHeader);
	}

	/**
	 * Step 1 of the AWS Signature version 4 calculation. Refer to http://docs.aws
	 * .amazon.com/general/latest/gr/sigv4-create-canonical-request.html to generate
	 * the canonical request.
	 */
	private String createCanonicalRequest(SignableRequest<?> request, String contentSha256) {
		final String strFormat = "%s\n%s\n%s\n%s\n%s\n%s";
		return String.format(strFormat, request.getHttpMethod(),
				getCanonicalizedResourcePath(request.getUri().getPath()), getCanonicalizedQueryString(request),
				getCanonicalizedHeaderString(request), getSignedHeadersString(request), contentSha256);
	}

	/**
	 * Step 2 of the AWS Signature version 4 calculation. Refer to http://docs.aws
	 * .amazon.com/general/latest/gr/sigv4-create-string-to-sign.html.
	 */
	private String createStringToSign(String canonicalRequest) {

		final String strFormat = "%s\n%s\n%s\n%s";

		return String.format(strFormat, AWS4_SIGNING_ALGORITHM, this.formattedSigningDateTime, this.scope,
				Hex.encodeHexString(hash(canonicalRequest)));
	}

	/**
	 * Step 3 of the AWS Signature version 4 calculation. It involves deriving the
	 * signing key and computing the signature. Refer to http://docs.aws.amazon
	 * .com/general/latest/gr/sigv4-calculate-signature.html
	 */
	private byte[] deriveSigningKey(String sk) {

		return newSigningKey(sk, this.formattedSigningDate, this.regionName, this.serviceName);
	}

	/**
	 * Step 3 of the AWS Signature version 4 calculation. It involves deriving the
	 * signing key and computing the signature. Refer to http://docs.aws.amazon
	 * .com/general/latest/gr/sigv4-calculate-signature.html
	 */
	private byte[] computeSignature(String stringToSign, byte[] signingKey) {
		try {
			return sign(stringToSign.getBytes("UTF-8"), signingKey, HMAC_SHA256_ALGORITHM);
		} catch (UnsupportedEncodingException e) {
			throw new KscClientException(e);
		}
	}

	/**
	 * Creates the authorization header to be included in the request.
	 */
	private String buildAuthorizationHeader(SignableRequest<?> request, byte[] signature, String ak) {

		final String credential = "Credential=" + ak + "/" + this.scope;
		final String signerHeaders = "SignedHeaders=" + getSignedHeadersString(request);
		final String signatureHeader = "Signature=" + Hex.encodeHexString(signature);

		return String.format("%s %s, %s, %s", AWS4_SIGNING_ALGORITHM, credential, signerHeaders, signatureHeader);
	}

	private String getCanonicalizedHeaderString(SignableRequest<?> request) {
		final List<String> sortedHeaders = new ArrayList<String>(request.getHeaders().keySet());
		Collections.sort(sortedHeaders, String.CASE_INSENSITIVE_ORDER);

		final Map<String, String> requestHeaders = request.getHeaders();
		StringBuilder buffer = new StringBuilder();
		for (String header : sortedHeaders) {
			String key = StringUtils.lowerCase(header);
			String value = requestHeaders.get(header);

			StringUtils.appendCompactedString(buffer, key);
			buffer.append(":");
			if (value != null) {
				StringUtils.appendCompactedString(buffer, value);
			}

			buffer.append("\n");
		}

		return buffer.toString();
	}

	private String getSignedHeadersString(SignableRequest<?> request) {
		final List<String> sortedHeaders = new ArrayList<String>(request.getHeaders().keySet());
		Collections.sort(sortedHeaders, String.CASE_INSENSITIVE_ORDER);

		StringBuilder buffer = new StringBuilder();
		for (String header : sortedHeaders) {
			if (buffer.length() > 0)
				buffer.append(";");
			buffer.append(StringUtils.lowerCase(header));
		}

		return buffer.toString();
	}

	/**
	 * Calculate the hash of the request's payload. Subclass could override this
	 * method to provide different values for "x-amz-content-sha256" header or do
	 * any other necessary set-ups on the request headers. (e.g. aws-chunked uses a
	 * pre-defined header value, and needs to change some headers relating to
	 * content-encoding and content-length.)
	 */
	private String calculateContentHash(SignableRequest<?> request) {

		try {
			InputStream payloadStream = request.getContent();
			if (payloadStream == null) {
				payloadStream = new ByteArrayInputStream(new byte[0]);
			}
			String contentSha256 = Hex.encodeHexString(hash(payloadStream));
			payloadStream.reset();

			return contentSha256;
		} catch (IOException e) {
			throw new KscClientException(e);
		}
	}

	/**
	 * Generates a new signing key from the given parameters and returns it.
	 */

	private byte[] newSigningKey(String sk, String dateStamp, String regionName, String serviceName) {
		try {
			byte[] kSecret = ("AWS4" + sk).getBytes("UTF-8");
			byte[] kDate = sign(dateStamp, kSecret, HMAC_SHA256_ALGORITHM);
			byte[] kRegion = sign(regionName, kDate, HMAC_SHA256_ALGORITHM);
			byte[] kService = sign(serviceName, kRegion, HMAC_SHA256_ALGORITHM);

			return sign(AWS4_TERMINATOR, kService, HMAC_SHA256_ALGORITHM);
		} catch (UnsupportedEncodingException e) {
			throw new KscClientException(e);
		}
	}

	private byte[] sign(String stringData, byte[] key, String algorithm) throws KscClientException {
		try {
			byte[] data = stringData.getBytes(UTF8);
			return sign(data, key, algorithm);
		} catch (Exception e) {
			throw new KscClientException("Unable to calculate a request signature: " + e.getMessage(), e);
		}
	}

	private byte[] sign(byte[] data, byte[] key, String algorithm) throws KscClientException {
		try {
			Mac mac = Mac.getInstance(algorithm);
			mac.init(new SecretKeySpec(key, algorithm));
			return mac.doFinal(data);
		} catch (Exception e) {
			throw new KscClientException("Unable to calculate a request signature: " + e.getMessage(), e);
		}
	}

	private static byte[] hash(String text) {
		try {
			MessageDigest md = MessageDigest.getInstance("SHA-256");
			md.update(text.getBytes(UTF8));
			return md.digest();
		} catch (Exception e) {
			throw new KscClientException("Unable to compute hash while signing request: " + e.getMessage(), e);
		}
	}

	private byte[] hash(InputStream input) throws KscClientException {
		try {
			MessageDigest md = MessageDigest.getInstance("SHA-256");
			DigestInputStream digestInputStream = new DigestInputStream(input, md);
			byte[] buffer = new byte[1024];
			while (digestInputStream.read(buffer) > -1)
				;
			return digestInputStream.getMessageDigest().digest();
		} catch (Exception e) {
			throw new KscClientException("Unable to compute hash while signing request: " + e.getMessage(), e);
		}
	}

	private String getCanonicalizedQueryString(SignableRequest<?> request) {
		/*
		 * If we're using POST and we don't have any request payload content, then any
		 * request query parameters will be sent as the payload, and not in the actual
		 * query string.
		 */
		if (SdkHttpUtils.usePayloadForQueryParameters(request))
			return "";
		return getCanonicalizedQueryString(request.getParameters());
	}

	/**
	 * Examines the specified query string parameters and returns a canonicalized
	 * form.
	 * <p>
	 * The canonicalized query string is formed by first sorting all the query
	 * string parameters, then URI encoding both the key and value and then joining
	 * them, in order, separating key value pairs with an '&'.
	 *
	 * @param parameters The query string parameters to be canonicalized.
	 *
	 * @return A canonicalized form for the specified query string parameters.
	 */
	private String getCanonicalizedQueryString(Map<String, String> parameters) {

		final SortedMap<String, String> sorted = new TreeMap<>();

		/**
		 * Signing protocol expects the param values also to be sorted after url
		 * encoding in addition to sorted parameter names.
		 */
		for (Map.Entry<String, String> entry : parameters.entrySet()) {
			final String encodedParamName = SdkHttpUtils.urlEncode(entry.getKey(), false);
			final String encodedValues = SdkHttpUtils.urlEncode(entry.getValue(), false);
			sorted.put(encodedParamName, encodedValues);
		}

		final StringBuilder result = new StringBuilder();
		for (Map.Entry<String, String> entry : sorted.entrySet()) {
			if (result.length() > 0) {
				result.append("&");
			}
			result.append(entry.getKey()).append("=").append(entry.getValue());
		}

		return result.toString();
	}

	private String getCanonicalizedResourcePath(String resourcePath) {
		if (resourcePath == null || resourcePath.isEmpty()) {
			return "/";
		}

		return resourcePath.startsWith("/") ? resourcePath : "/".concat(resourcePath);
	}
}
